import pathlib

from qdrant_client import QdrantClient
from qdrant_client.http import models
from sentence_transformers import SentenceTransformer
import torch
from pathlib import Path

# 输入文本
texts = ["""
1．太阳之为病，脉浮、头项强痛而恶寒。
2．太阳病，发热、汗出、恶风、脉缓者，名为中风。
3．太阳病，或已发热，或未发热，必恶寒、体痛、呕逆、脉阴阳俱紧者，名为伤寒。
4．伤寒一日，太阳受之。脉若静者，为不传；颇欲吐，若躁烦，脉数急者，为传也。
5．伤寒二三日，阳明、少阳证不见者，为不传也。
6．太阳病，发热而渴，不恶寒者，为温病。若发汗已，身灼热者，名风温。风温为病，脉阴阳俱浮、自汗出、身重、多眠睡、鼻息必鼾、语言难出；若被下者，小便不利、直视失溲；若被火者，微发黄色，剧则如惊痫，时螈 ；若火熏之，一逆尚引日，再逆促命期。
7．病有发热恶寒者；发于阳也；无热恶寒者，发于阴也。发于阳，七日愈；发于阴，六日愈。以阳数七、阴数六故也。
8．太阳病，头痛至七日以上自愈者，以行其经尽故也。若欲作再经者，针足阳明，使经不传则愈。
9．太阳病欲解时，从巳至未上。
10．风家，表解而不了了者，十二日愈。
11．病患身大热，反欲得衣者，热在皮肤，寒在骨髓也；身大寒，反不欲近衣者，寒在皮肤，热在骨髓也。
12．太阳中风，阳浮而阴弱，阳浮者，热自发；阴弱者，汗自出。啬啬恶寒，淅淅恶风，翕翕发热， 鼻鸣干呕者，桂枝汤主之。
""", """桂枝汤方:桂枝（去皮，三两） 芍药（三两） 甘草（炙，二两） 生姜（切，三两） 大枣（擘，十二枚）
上五味， 咀三味，以水七升，微火煮取三升，去滓，适寒温，服一升。服已须臾，啜热稀粥一升余，以助药力，温覆令一时许，遍身 微似有汗者益佳；不可令如水流离，病必不除。
若一服汗出病瘥，停后服，不必尽剂；若不汗，更服，根据前法；又不汗，后服小促其间，半日许令三服尽。若病重者，一日一夜服，周时观之，服一剂尽，病证犹在者，更作服；若汗不出，乃服至二、三剂。禁生冷、粘滑、肉面、五辛、酒酪、臭恶等物。
13．太阳病，头痛、发热、汗出、恶风，桂枝汤主之。方二。（用前第一方。）
14．太阳病，项背强KT KT ，反汗出恶风者，桂枝加葛根汤主之。
葛根汤方:葛根（四两） 麻黄（去节，三两） 芍药（二两） 生姜（切，三两） 甘草（炙，二两） 大枣（擘，十二枚） 桂枝（去皮，二两）
上七味；以水一斗，先煮麻黄、葛根，减二升，去上沫，内诸药，煮取三升，去滓，温服一升。覆取微似汗，不须啜粥，余如桂枝法将息及禁忌。
""", """(臣亿等谨按仲景本论，太阳中风自汗用桂枝，伤寒无汗用麻黄，今证云汗出恶风，而方中有麻黄，恐非本意也。第三卷有葛根汤证云，无汗恶风，正与此方同，是合用麻黄也。此云桂枝加葛根汤，恐是桂枝中但加葛根耳。)
15．太阳病，下之后，其气上冲者，可与桂枝汤，方用前法；若不上冲者，不得与之。方四。
16．太阳病三日，已发汗，若吐、若下、若温针，仍不解者，此为坏病，桂枝不中与之也。观其脉证，知犯何逆，随证治之。桂枝本为解肌，若其人脉浮紧、发热、汗不出者，不可与之也。常须识此，勿令误也。方五。
17．若酒客病，不可与桂枝汤，得之则呕，以酒客不喜甘故也。
18．喘家，作桂枝汤，加浓朴、杏子佳。方六。
19．凡服桂枝汤吐者，其后必吐脓血也。
20．太阳病，发汗，遂漏不止，其人恶风，小便难，四肢微急，难以屈伸者，桂枝加附子汤主之。方七。
桂枝（去皮，三两） 芍药（三两） 甘草（炙，三两） 生姜（切，三两） 大枣（擘，十二枚） 附子（一枚，炮，去皮，破八片）
上六味，以水七升，煮取三升，去滓，温服一升。本云桂枝汤，今加附子，将息如前法。
21．太阳病，下之后，脉促、胸满者，桂枝去芍药汤主之。方八。（促，一作纵。）
"""]



# 初始化Qdrant客户端（假设本地运行）
client = QdrantClient(host="192.168.2.230", port=6333)

# 加载中文预训练模型
models_dir = Path.joinpath(Path.cwd().parent.parent.parent, "local_models", "bge-small-zh-v1.5")
print("models_dir:", models_dir)
model_path = Path.joinpath(models_dir, "model.safetensors")
print("model_path:", model_path)
if not Path.exists(model_path):
    print("models/bge-small-zh-v1.5/model.safetensors 文件不存在，请检查！")
# 使用在线模型 ID 加载模型
model = SentenceTransformer("BAAI/bge-small-zh-v1.5")

# 或者使用本地路径加载模型
#model = SentenceTransformer("E:\workspace\python\pear-admin-flask\models\bge-small-zh-v1.5\model.safetensors")
embeddings = model.encode(texts)

# 创建新集合，指定向量维度和使用的距离度量方式
if not client.collection_exists("chinese_medicine"):
    client.create_collection("chinese_medicine",
        vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE))

for i, text in enumerate(texts):
    embedding = embeddings[i].tolist()

    # 创建条目对象
    point_id = i + 1
    payload = {"text": text}
    print("id:", point_id, "vector:", embedding, "vector.len:", len(embedding), "payload:", payload)
    # 插入数据到 Qdrant
    # client.upsert(
    #     collection_name="chinese_medicine",
    #     points=[
    #         models.PointStruct(
    #             id=point_id,
    #             payload=payload,
    #             vector=embedding
    #         )
    #     ]
    # )

# 检索相似文本
query_text = "太阳中风"
query_embedding = model.encode(query_text)

hits = client.search(
    collection_name="chinese_medicine",
    query_vector=query_embedding.tolist(),
    limit=1
)

results = [hit.payload["text"] for hit in hits]

print("最相似的文本片段：")
print("\n".join(results))